#encoding: utf8
from scapy.all import *
import scapy_http.http
import time
import numpy as np
from enum import Enum
class PacketStatus(Enum):
    Req = 0
    Rsp = 1
    Others = 2


def http_header_string_to_dict(http_header_str):
    http_header_str = http_header_str.strip()
    http_head_dict = {}
    http_head_list = http_header_str.split("\r\n")[1:]
    for http_head in http_head_list:
        d = http_head.split(":")
        # Date: Thu, 26 Apr 2018 08:29:22 GMT
        head_k = d[0].strip()
        head_v = ":".join(d[1:]).strip()
        http_head_dict[head_k] = head_v

    return http_head_dict

def http_header_content_type_table(content_type):

    pass

if __name__ == '__main__':
    all_packets = rdpcap("image_http_request.pcap")

    sessions = all_packets.sessions()

    # 查看每个session的包大小
    # tcp_list_len = []
    # for k,v in sessions.items():
    #     tcp_list_len.append(len(v))
    # tcp_list_len.sort()
    # print tcp_list_len

    # TCP 125.216.241.118:3008 > 192.168.227.131:33777  http response
    # TCP 192.168.227.131:33777 > 125.216.241.118:3008  http request

    # 查看那些包不带数据
    # rsp = sessions["TCP 125.216.241.118:3008 > 192.168.227.131:33777"]
    # idx = 0
    # idx_list = []
    # for rr in rsp:
    #     idx += 1
    #     if rr.len == 40:
    #         idx_list.append(idx)

    
    # 抓取http request and response
    status = PacketStatus.Others
    is_end = True


    for s_key, s_val in sessions.items():
        header_buf = ""
        content_buf = None
        data_buf_base_seq = -1
        http_head_dict = None

        for pkt in s_val:

            if pkt.haslayer("Raw"):

                data = pkt.getlayer("Raw").fields["load"]

                if status == PacketStatus.Others:
                    header_buf = ""

                    if data.find("HTTP/1.1 200 OK") == 0:
                        status = PacketStatus.Rsp
                        content_buf = None
                        data_buf_base_seq = pkt.seq

                    elif data.find("GET") == 0 or data.find("POST") == 0:
                        status = PacketStatus.Req

                    else:
                        status =  PacketStatus.Others
                        continue


                
                http_header_end_idx = header_buf.find("\r\n\r\n")
                
                if http_header_end_idx == -1:
                    header_buf += data
                    continue

                http_header = header_buf[:http_header_end_idx]
                http_head_dict = http_header_string_to_dict(http_header)

                # http request 遇到两个\r\n就是请求头结束
                if status == PacketStatus.Req:
                    is_end = True
                    continue

                # 处理 http body
                content_size = int(http_head_dict["Content-Length"])
                if content_buf == None:
                    content_buf = ['\x00'] * content_size
                    data_len = len(header_buf[http_header_end_idx + 4:])
                    content_buf[:data_len] = header_buf[http_header_end_idx + 4:]

                
                offset = pkt.seq - data_buf_base_seq - http_header_end_idx - 4
                data_len = len(data)
                content_buf[offset: offset+data_len] = data

                # http response数据接受完成
                if status == PacketStatus.Rsp and \
                        offset + data_len < content_size:
                    continue

                print(http_header + "\n")

                is_end = True
                # 只处理图片

                content_type = http_head_dict["Content-Type"]
                if content_type.find("image") < 0:
                    continue
                img_ext = content_type.split("/")[-1]
                time_stamp = pkt.time
                
                if not os.path.exists("image/"):
                    os.makedirs("image/")

                f = open("image/%s.%s" %(time_stamp, img_ext), "wb")
                content_buf = "".join(content_buf)
                f.write(content_buf)
                f.close()
                print "rsp receive complete"


            else:

                status = PacketStatus.Others
                if is_end == False:
                    print "error"

                is_end = True

